import ast
from collections import OrderedDict
import logging
import astunparse
from functools import partial


class Required:
    pass


class UnRequired:
    pass


class API:
    def __init__(self, name: str, params: OrderedDict):
        self.name = name
        self.params = OrderedDict()

        def to_str(v):
            if v is Required:
                return v
            elif isinstance(v, str):
                return "'{}'".format(v)
            else:
                return str(v)

        for k, v in params.items():
            self.params[k] = to_str(v)

    def parse_args(self, call_name: str, args_str: str):
        # expr必须满足(**)的形式
        assert len(args_str) >= 2 and args_str[0] == "(" and args_str[-1] == ")"

        try:
            ast_node = ast.parse("whatever_call_name" + args_str)
            try:
                call_node = ast_node.body[0].value
                assert isinstance(call_node, ast.Call)
            except:
                raise ValueError("ast_node.body[0].value is not ast.Call")
        except:
            raise ValueError("can't parse code:\n{}".format(args_str))

        # 将所有实参看成一个参数
        if len(self.params) == 1:
            k = list(self.params.keys())[0]
            if k.startswith('*'):
                v = args_str[1:-1]  # remove ()
                return OrderedDict([(k, v), ("call_name", call_name)])

        args = OrderedDict()

        # 没有指定名字的参数
        param_iter = iter(self.params.keys())
        assert len(call_node.args) <= len(self.params)
        for arg in call_node.args:
            if isinstance(arg, ast.Starred):
                logging.debug("find *{}".format(arg.value.id))
                args['*'] = arg.value.id
            else:
                args[next(param_iter)] = astunparse.unparse(arg).strip()  # remove \n

        # 指定名字的参数
        for keyword in call_node.keywords:
            if keyword.arg is None:  # TODO
                logging.debug("find **{}".format(keyword.value.id))
                args['**'] = keyword.value.id
            else:
                assert keyword.arg in self.params
                assert keyword.arg not in args
                args[keyword.arg] = astunparse.unparse(keyword.value).strip()  # remove \n

        args["call_name"] = call_name
        return args


class APIPt(API):
    pass


class APIMs(API):
    def __init__(self, name: str, params: OrderedDict, P_attrs=None):
        self.is_Primitive = name.startswith('P.')
        if self.is_Primitive:
            self.P_attrs = P_attrs if P_attrs else set()
        super(APIMs, self).__init__(name, params)

    def create_args(self, params_pt, args_pt, ms2pt_map, explicit_map):
        args = OrderedDict()

        #  遍历mindspore的参数表
        for k in self.params.keys():
            if k in ms2pt_map:  # 有pytorch对应的映射参数
                if ms2pt_map[k] in args_pt:  # 用户给pytorch的api赋值了
                    args[k] = args_pt[ms2pt_map[k]]
                elif self.params[k] != params_pt[ms2pt_map[k]]:  # 用户没有赋值，但两个框架的初始值不一样
                    args[k] = params_pt[ms2pt_map[k]]
            else:  # 没有pytorch对应的映射参数
                if k in explicit_map:  # 强行显示的参数
                    args[k] = explicit_map[k]
                elif self.params[k] is Required:
                    args[k] = "<Required>"

        # 在pytorch实参表里发现*或**
        for star in ('*', '**'):
            if star in args_pt:
                args[star] = args_pt[star]

        return args


class MappingHelper:
    def __init__(self, ms_api: APIMs, pt_api: APIPt, ms2pt_mapping=None, gen_explicit_map=None, export_key=None):
        if ms2pt_mapping is None:
            ms2pt_mapping = {}
        if gen_explicit_map is None:
            gen_explicit_map = lambda params_pt, args_pt: {}
        self.ms_api = ms_api
        self.pt_api = pt_api
        self.ms2pt_mapping = ms2pt_mapping
        self.gen_explicit_map = gen_explicit_map
        if export_key is not None:
            self.export_key = export_key
        else:
            self.export_key = not ms_api.is_Primitive

    def gen_args_expr(self, args):
        expr = ''
        for k, v in args.items():
            if expr:
                expr += ', '
            sym = '' if k in ('*', '**') else '='
            if self.export_key:
                expr += k + sym
            expr += v
        return expr

    def gen_args_expr_for_P(self, args, P_attrs):
        args_attrs = OrderedDict([(k, v) for k, v in args.items() if k in P_attrs])
        args_ios = OrderedDict([(k, v) for k, v in args.items() if k not in P_attrs])
        return self.gen_args_expr(args_attrs), self.gen_args_expr(args_ios)

    def convert(self, call_name_pt: str, args_str_pt: str):
        # args_pt的value值都是字符串
        args_pt = self.pt_api.parse_args(call_name_pt, args_str_pt)

        # args_ms的value值都是字符串
        explicit_map = self.gen_explicit_map(self.pt_api.params, args_pt)
        args_ms = self.ms_api.create_args(self.pt_api.params, args_pt, self.ms2pt_mapping, explicit_map)

        if self.ms_api.is_Primitive:
            if self.pt_api.name == '.size' and 'idx' in args_pt:
                args_expr = self.gen_args_expr(args_ms)
                expr_ms = "%s()(%s)[%s]" % (self.ms_api.name, args_expr, args_pt['idx'])
            else:
                expr_attrs, expr_ios = self.gen_args_expr_for_P(args_ms, self.ms_api.P_attrs)
                expr_ms = "%s(%s)(%s)" % (self.ms_api.name, expr_attrs, expr_ios)
        else:
            ms_expr = self.gen_args_expr(args_ms)
            expr_ms = "%s(%s)" % (self.ms_api.name, ms_expr)
        return expr_ms


def gen_explicit_map_nnSequential(params_pt, args_pt):
    args = args_pt['*args']
    return {"*args": "[{}]".format(args)}


def gen_explicit_map_nnMaxPool2d(params_pt, args_pt):
    # TODO warning
    if 'padding' in args_pt:
        padding = args_pt['padding']
    else:
        padding = params_pt['padding']
    if padding.strip() in ("0", "(0,0)", "(0, 0)"):
        pad_mode = "'valid'"
    else:
        pad_mode = "'same'"
    return {"pad_mode": pad_mode}


def gen_explicit_map_F_max_pool2d(params_pt, args_pt):
    # TODO warning
    if 'padding' in args_pt:
        padding = args_pt['padding']
    else:
        padding = params_pt['padding']
    if padding.strip() in ("0", "(0,0)", "(0, 0)"):
        padding = "'valid'"
    else:
        padding = "'same'"
    return {"padding": padding}


def gen_explicit_map_one_delta(params_pt, args_pt, k_ms, k_pt):
    v = args_pt[k_pt] if k_pt in args_pt else params_pt[k_pt]
    v = v.strip()

    def is_number(s):
        try:
            float(s)
            return True
        except ValueError:
            return False

    if is_number(v):
        return {k_ms: str(1 - float(v))}
    else:
        return {k_ms: "1.0 - " + v}


# ---------------------------- 映射表 ----------------------------
NN_MAPPING = {
    'nn.Sequential': MappingHelper(**{"ms_api": APIMs('nn.SequentialCell', OrderedDict([('*args', Required)])),
                                      "pt_api": APIPt('nn.Sequential', OrderedDict([('*args', Required)])),
                                      "gen_explicit_map": gen_explicit_map_nnSequential,
                                      "export_key": False
                                      }),
    'nn.Conv2d': MappingHelper(**{"ms_api": APIMs('nn.Conv2d', OrderedDict(in_channels=Required,
                                                                           out_channels=Required,
                                                                           kernel_size=Required,
                                                                           stride=1,
                                                                           pad_mode='same',
                                                                           padding=0,
                                                                           dilation=1,
                                                                           group=1,
                                                                           has_bias=False,
                                                                           weight_init='normal',
                                                                           bias_init='zeros')),
                                  "pt_api": APIPt('nn.Conv2d', OrderedDict(in_channels=Required,
                                                                           out_channels=Required,
                                                                           kernel_size=Required,
                                                                           stride=1,
                                                                           padding=0,
                                                                           dilation=1,
                                                                           groups=1,
                                                                           bias=True,
                                                                           padding_mode='zeros')),
                                  "ms2pt_mapping": {'in_channels': 'in_channels',
                                                    'out_channels': 'out_channels',
                                                    'kernel_size': 'kernel_size',
                                                    'stride': 'stride',
                                                    'padding': 'padding',
                                                    'dilation': 'dilation',
                                                    'group': 'groups',
                                                    'has_bias': 'bias'},
                                  "gen_explicit_map": (lambda params_pt, args_pt: {"pad_mode": "'pad'"})
                                  }),
    'nn.BatchNorm2d': MappingHelper(**{"ms_api": APIMs('nn.BatchNorm2d', OrderedDict(num_features=Required,
                                                                                     eps=1e-5,
                                                                                     momentum=0.9,
                                                                                     affine=True,
                                                                                     gamma_init='ones',
                                                                                     beta_init='zeros',
                                                                                     moving_mean_init='zeros',
                                                                                     moving_var_init='ones',
                                                                                     use_batch_statistics=True)),
                                       "pt_api": APIPt('nn.BatchNorm2d', OrderedDict(num_features=Required,
                                                                                     eps=1e-5,
                                                                                     momentum=0.1,
                                                                                     affine=True,
                                                                                     track_running_stats=True)),
                                       "ms2pt_mapping": {"num_features": "num_features",
                                                         "eps": "eps",
                                                         "affine": "affine",
                                                         "use_batch_statistics": "track_running_stats"},
                                       "gen_explicit_map": partial(gen_explicit_map_one_delta,
                                                                   k_ms="momentum", k_pt="momentum")
                                       }),
    'nn.ReLU': MappingHelper(**{"ms_api": APIMs('nn.ReLU', OrderedDict()),
                                "pt_api": APIPt('nn.ReLU', OrderedDict(inplace=False)),
                                "ms2pt_mapping": {}}),
    'nn.ReLU6': MappingHelper(**{"ms_api": APIMs('nn.ReLU6', OrderedDict()),
                                 "pt_api": APIPt('nn.ReLU6', OrderedDict(inplace=False)),
                                 "ms2pt_mapping": {}}),
    'nn.Linear': MappingHelper(**{"ms_api": APIMs('nn.Dense', OrderedDict(in_channels=Required,
                                                                          out_channels=Required,
                                                                          weight_init='normal',
                                                                          bias_init='zeros',
                                                                          has_bias=True,
                                                                          activation=None)),
                                  "pt_api": APIPt('nn.Linear', OrderedDict(in_features=Required,
                                                                           out_features=Required,
                                                                           bias=True)),
                                  "ms2pt_mapping": {"in_channels": "in_features",
                                                    "out_channels": "out_features",
                                                    "has_bias": "bias"}
                                  }),
    'nn.MaxPool2d': MappingHelper(**{"ms_api": APIMs('nn.MaxPool2d', OrderedDict(kernel_size=1,
                                                                                 stride=1,
                                                                                 pad_mode="valid")),
                                     "pt_api": APIPt('nn.MaxPool2d', OrderedDict(kernel_size=Required,
                                                                                 stride=None,
                                                                                 padding=0,
                                                                                 dilation=1,
                                                                                 return_indices=False,
                                                                                 ceil_mode="False")),
                                     "ms2pt_mapping": {"kernel_size": "kernel_size",
                                                       "stride": "stride"},
                                     "gen_explicit_map": gen_explicit_map_nnMaxPool2d
                                     }),
    'nn.AvgPool2d': MappingHelper(**{"ms_api": APIMs('nn.AvgPool2d', OrderedDict(kernel_size=1,
                                                                                 stride=1,
                                                                                 pad_mode="valid")),
                                     "pt_api": APIPt('nn.AvgPool2d', OrderedDict(kernel_size=Required,
                                                                                 stride=None,
                                                                                 padding=0,
                                                                                 dilation=1,
                                                                                 return_indices=False,
                                                                                 ceil_mode="False")),
                                     "ms2pt_mapping": {"kernel_size": "kernel_size",
                                                       "stride": "stride"},
                                     "gen_explicit_map": gen_explicit_map_nnMaxPool2d
                                     }),
    'nn.Dropout': MappingHelper(**{"ms_api": APIMs('nn.Dropout', OrderedDict(keep_prob=0.5,
                                                                             seed0=0,
                                                                             seed1=0,
                                                                             dtype="mstype.float32")),  # TODO
                                   "pt_api": APIPt('nn.Dropout', OrderedDict(p=0.5,
                                                                             inplace=False)),
                                   "ms2pt_mapping": {"keep_prob": "p"},
                                   "gen_explicit_map": partial(gen_explicit_map_one_delta,
                                                               k_ms="keep_prob", k_pt="p")
                                   }),

    # 'nn.AdaptiveAvgPool2d': MappingHelper(**{"ms_api": APIMs('P.ReduceMean', OrderedDict(keep_dims=False,
    #                                                                                      x=Required,
    #                                                                                      axis=())),
    #                                          "pt_api": APIPt('nn.AdaptiveAvgPool2d', OrderedDict(output_size=Required)),
    #                                          "ms2pt_mapping": {"kernel_size": "kernel_size",
    #                                                            "stride": "stride"},
    #                                          "gen_explicit_map": gen_explicit_map_nnMaxPool2d
    #                                          }),
}
# 设置别名 nn. = torch.nn.
NN_MAPPING.update(dict([("torch." + k, v) for k, v in NN_MAPPING.items()]))

F_MAPPING = {
    'F.relu': MappingHelper(**{"ms_api": APIMs('P.ReLU', OrderedDict(input=Required)),
                               "pt_api": APIPt('F.relu', OrderedDict(input=Required, inplace=False)),
                               "ms2pt_mapping": {"input": "input"},
                               }),
    'F.relu6': MappingHelper(**{"ms_api": APIMs('P.ReLU6', OrderedDict(input=Required)),
                                "pt_api": APIPt('F.relu6', OrderedDict(input=Required, inplace=False)),
                                "ms2pt_mapping": {"input": "input"},
                                }),
    'F.max_pool2d': MappingHelper(**{"ms_api": APIMs('P.MaxPool', OrderedDict(ksize=1,
                                                                              strides=1,
                                                                              padding="valid",
                                                                              input=Required),
                                                     P_attrs={"ksize", "strides", "padding"}),
                                     "pt_api": APIPt('F.max_pool2d', OrderedDict(input=Required,
                                                                                 kernel_size=Required,
                                                                                 stride=None,
                                                                                 padding=0,
                                                                                 dilation=1,
                                                                                 ceil_mode=False,
                                                                                 return_indices=False)),
                                     "ms2pt_mapping": {"ksize": "kernel_size",
                                                       "strides": "stride",
                                                       "input": "input",
                                                       },
                                     "gen_explicit_map": gen_explicit_map_F_max_pool2d
                                     }),
    'F.avg_pool2d': MappingHelper(**{"ms_api": APIMs('P.AvgPool', OrderedDict(ksize=1,
                                                                              strides=1,
                                                                              padding="valid",
                                                                              input=Required),
                                                     P_attrs={"ksize", "strides", "padding"}),
                                     "pt_api": APIPt('F.avg_pool2d', OrderedDict(input=Required,
                                                                                 kernel_size=Required,
                                                                                 stride=None,
                                                                                 padding=0,
                                                                                 dilation=1,
                                                                                 ceil_mode=False,
                                                                                 return_indices=False)),
                                     "ms2pt_mapping": {"ksize": "kernel_size",
                                                       "strides": "stride",
                                                       "input": "input",
                                                       },
                                     "gen_explicit_map": gen_explicit_map_F_max_pool2d
                                     }),
}
# 设置别名 F = nn.functional = torch.nn.functional
F_MAPPING.update({
    **dict([("nn.functional." + k[2:], v) for k, v in F_MAPPING.items()]),
    **dict([("torch.nn.functional." + k[2:], v) for k, v in F_MAPPING.items()])
})

TORCH_DOT_MAPPING = {
    'torch.flatten': MappingHelper(**{"ms_api": APIMs('P.Flatten', OrderedDict(input=Required)),
                                      "pt_api": APIPt('torch.flatten', OrderedDict(input=Required,
                                                                                   start_dim=0,
                                                                                   end_dim=-1)),
                                      "ms2pt_mapping": {"input": "input"}
                                      }),
    'torch.cat': MappingHelper(**{"ms_api": APIMs('P.Concat',
                                                  OrderedDict(axis=0, input=Required),
                                                  P_attrs={"axis"}),
                                  "pt_api": APIPt('torch.flatten', OrderedDict(tensors=Required, dim=0, out=None)),
                                  "ms2pt_mapping": {"input": "tensors",
                                                    "axis": "dim"}
                                  }),
}

TENSOR_DOT_MAPPING = {
    # https://pytorch.org/docs/stable/tensors.html#torch.Tensor.view
    '.view': MappingHelper(**{"ms_api": APIMs('P.Reshape', OrderedDict(x=Required, shape=Required)),
                              "pt_api": APIPt('.view', OrderedDict([('*shape', Required)])),
                              "ms2pt_mapping": {"x": "call_name"},
                              "gen_explicit_map": (lambda params_pt, args_pt: {"shape": "(" + args_pt["*shape"] + ",)"})
                              }),
    '.size': MappingHelper(**{"ms_api": APIMs('P.Shape', OrderedDict(x=Required)),
                              "pt_api": APIPt('.size', OrderedDict([('idx', Required)])),
                              "ms2pt_mapping": {"x": "call_name"}
                              }),
    '.flatten': MappingHelper(**{"ms_api": APIMs('P.Flatten', OrderedDict(input=Required)),
                                 "pt_api": APIPt('.flatten', OrderedDict(start_dim=0,
                                                                         end_dim=-1)),
                                 "ms2pt_mapping": {"input": "call_name"}
                                 }),
    '.reshape': MappingHelper(**{"ms_api": APIMs('P.Reshape', OrderedDict(x=Required, shape=Required)),
                                 "pt_api": APIPt('.reshape', OrderedDict([('*shape', Required)])),
                                 "ms2pt_mapping": {"x": "call_name"},
                                 "gen_explicit_map": (
                                     lambda params_pt, args_pt: {"shape": "(" + args_pt["*shape"] + ",)"})
                                 }),
    '.mean': MappingHelper(**{"ms_api": APIMs('P.ReduceMean', OrderedDict(keep_dims=False,
                                                                          input=Required,
                                                                          axis=()),
                                              P_attrs={"keep_dims"}),
                              "pt_api": APIPt('.mean', OrderedDict(dim=None,
                                                                   keepdim=False)),
                              "ms2pt_mapping": {"keep_dims": "keepdim",
                                                "axis": "dim",
                                                "input": "call_name"},
                              }),
    '.squeeze': MappingHelper(**{"ms_api": APIMs('P.ReduceMean', OrderedDict(input=Required,
                                                                             axis=()),
                                                 P_attrs={"axis"}),
                                 "pt_api": APIPt('.squeeze', OrderedDict(dim=None)),
                                 "ms2pt_mapping": {"axis": "dim",
                                                   "input": "call_name"},
                                 }),
}

ALL_MAPPING = {**NN_MAPPING, **F_MAPPING, **TORCH_DOT_MAPPING, **TENSOR_DOT_MAPPING}

# ---------------------------- 支持/不支持的api列表 ----------------------------
NN_LIST = [
    "nn.Module",
    "nn.CELU",
    "nn.ELU",
    "nn.GLU",
    "nn.GELU",
    "nn.Hardshrink",
    "nn.Hardtanh",
    "nn.LeakyReLU",
    "nn.LogSigmoid",
    "nn.LogSoftmax",
    "nn.PReLU",
    "nn.RReLU",
    "nn.ReLU",
    "nn.ReLU6",
    "nn.SELU",
    "nn.Sigmoid",
    "nn.Softmax",
    "nn.Softmax2d",
    "nn.Softmin",
    "nn.Softplus",
    "nn.Softshrink",
    "nn.Softsign",
    "nn.Tanh",
    "nn.Tanhshrink",
    "nn.Threshold",
    "nn.MultiheadAttention",
    "nn.AdaptiveLogSoftmaxWithLoss",
    "nn.BatchNorm1d",
    "nn.BatchNorm2d",
    "nn.BatchNorm3d",
    "nn.SyncBatchNorm",
    "nn.Container",
    "nn.ModuleDict",
    "nn.ModuleList",
    "nn.ParameterDict",
    "nn.ParameterList",
    "nn.Sequential",
    "nn.Conv1d",
    "nn.Conv2d",
    "nn.Conv3d",
    "nn.ConvTranspose1d",
    "nn.ConvTranspose2d",
    "nn.ConvTranspose3d",
    "nn.CosineSimilarity",
    "nn.PairwiseDistance",
    "nn.AlphaDropout",
    "nn.Dropout",
    "nn.Dropout2d",
    "nn.Dropout3d",
    "nn.FeatureAlphaDropout",
    "nn.Fold",
    "nn.Unfold",
    "nn.InstanceNorm1d",
    "nn.InstanceNorm2d",
    "nn.InstanceNorm3d",
    "nn.Bilinear",
    "nn.Identity",
    "nn.Linear",
    "nn.BCELoss",
    "nn.BCEWithLogitsLoss",
    "nn.CTCLoss",
    "nn.CosineEmbeddingLoss",
    "nn.CrossEntropyLoss",
    "nn.HingeEmbeddingLoss",
    "nn.KLDivLoss",
    "nn.L1Loss",
    "nn.MSELoss",
    "nn.MarginRankingLoss",
    "nn.MultiLabelMarginLoss",
    "nn.MultiLabelSoftMarginLoss",
    "nn.MultiMarginLoss",
    "nn.NLLLoss",
    "nn.NLLLoss2d",
    "nn.PoissonNLLLoss",
    "nn.SmoothL1Loss",
    "nn.SoftMarginLoss",
    "nn.TripletMarginLoss",
    "nn.Module",
    "nn.CrossMapLRN2d",
    "nn.GroupNorm",
    "nn.LayerNorm",
    "nn.LocalResponseNorm",
    "nn.ConstantPad1d",
    "nn.ConstantPad2d",
    "nn.ConstantPad3d",
    "nn.ReflectionPad1d",
    "nn.ReflectionPad2d",
    "nn.ReplicationPad1d",
    "nn.ReplicationPad2d",
    "nn.ReplicationPad3d",
    "nn.ZeroPad2d",
    "nn.PixelShuffle",
    "nn.AdaptiveAvgPool1d",
    "nn.AdaptiveAvgPool2d",
    "nn.AdaptiveAvgPool3d",
    "nn.AdaptiveMaxPool1d",
    "nn.AdaptiveMaxPool2d",
    "nn.AdaptiveMaxPool3d",
    "nn.AvgPool1d",
    "nn.AvgPool2d",
    "nn.AvgPool3d",
    "nn.FractionalMaxPool2d",
    "nn.FractionalMaxPool3d",
    "nn.LPPool1d",
    "nn.LPPool2d",
    "nn.MaxPool1d",
    "nn.MaxPool2d",
    "nn.MaxPool3d",
    "nn.MaxUnpool1d",
    "nn.MaxUnpool2d",
    "nn.MaxUnpool3d",
    "nn.GRU",
    "nn.GRUCell",
    "nn.LSTM",
    "nn.LSTMCell",
    "nn.RNN",
    "nn.RNNBase",
    "nn.RNNCell",
    "nn.RNNCellBase",
    "nn.Embedding",
    "nn.EmbeddingBag",
    "nn.Upsample",
    "nn.UpsamplingBilinear2d",
    "nn.UpsamplingNearest2d",
    "nn.Transformer",
    "nn.TransformerEncoder",
    "nn.TransformerDecoder",
    "nn.TransformerEncoderLayer",
    "nn.TransformerDecoderLayer",

    # from .modules import *
    # from .parameter import Parameter as Parameter
    # from .parallel import DataParallel as DataParallel
    # from . import init as init
    # from . import utils as utils
    # from . import functional as functional
    # from . import parallel as parallel
    "nn.Parameter",
    "nn.DataParallel",
]
# 设置别名 nn. = torch.nn.
NN_LIST += ["torch." + name for name in NN_LIST]
NN_SUPPORTED = [x for x in NN_LIST if x in ALL_MAPPING]
NN_UNSUPPORTED = [x for x in NN_LIST if x not in ALL_MAPPING]

F_LIST = [
    "torch.nn.functional.conv1d",
    "torch.nn.functional.conv2d",
    "torch.nn.functional.conv3d",
    "torch.nn.functional.conv_transpose1d",
    "torch.nn.functional.conv_transpose2d",
    "torch.nn.functional.conv_transpose3d",
    "torch.nn.functional.unfold",
    "torch.nn.functional.fold",
    "torch.nn.functional.avg_pool1d",
    "torch.nn.functional.avg_pool2d",
    "torch.nn.functional.avg_pool3d",
    "torch.nn.functional.max_pool1d",
    "torch.nn.functional.max_pool2d",
    "torch.nn.functional.max_pool3d",
    "torch.nn.functional.max_unpool1d",
    "torch.nn.functional.max_unpool2d",
    "torch.nn.functional.max_unpool3d",
    "torch.nn.functional.lp_pool1d",
    "torch.nn.functional.lp_pool2d",
    "torch.nn.functional.adaptive_max_pool1d",
    "torch.nn.functional.adaptive_max_pool2d",
    "torch.nn.functional.adaptive_max_pool3d",
    "torch.nn.functional.adaptive_avg_pool1d",
    "torch.nn.functional.adaptive_avg_pool2d",
    "torch.nn.functional.adaptive_avg_pool3d",
    "torch.nn.functional.threshold",
    "torch.nn.functional.threshold_",
    "torch.nn.functional.relu",
    "torch.nn.functional.relu_",
    "torch.nn.functional.hardtanh",
    "torch.nn.functional.hardtanh_",
    "torch.nn.functional.relu6",
    "torch.nn.functional.elu",
    "torch.nn.functional.elu_",
    "torch.nn.functional.selu",
    "torch.nn.functional.celu",
    "torch.nn.functional.leaky_relu",
    "torch.nn.functional.leaky_relu_",
    "torch.nn.functional.prelu",
    "torch.nn.functional.rrelu",
    "torch.nn.functional.rrelu_",
    "torch.nn.functional.glu",
    "torch.nn.functional.gelu",
    "torch.nn.functional.logsigmoid",
    "torch.nn.functional.hardshrink",
    "torch.nn.functional.tanhshrink",
    "torch.nn.functional.softsign",
    "torch.nn.functional.softplus",
    "torch.nn.functional.softmin",
    "torch.nn.functional.softmax",
    "torch.nn.functional.softshrink",
    "torch.nn.functional.gumbel_softmax",
    "torch.nn.functional.log_softmax",
    "torch.nn.functional.tanh",
    "torch.nn.functional.sigmoid",
    "torch.nn.functional.batch_norm",
    "torch.nn.functional.instance_norm",
    "torch.nn.functional.layer_norm",
    "torch.nn.functional.local_response_norm",
    "torch.nn.functional.normalize",
    "torch.nn.functional.linear",
    "torch.nn.functional.bilinear",
    "torch.nn.functional.dropout",
    "torch.nn.functional.alpha_dropout",
    "torch.nn.functional.dropout2d",
    "torch.nn.functional.dropout3d",
    "torch.nn.functional.embedding",
    "torch.nn.functional.embedding_bag",
    "torch.nn.functional.one_hot",
    "torch.nn.functional.pairwise_distance",
    "torch.nn.functional.cosine_similarity",
    "torch.nn.functional.pdist",
    "torch.nn.functional.binary_cross_entropy",
    "torch.nn.functional.binary_cross_entropy_with_logits",
    "torch.nn.functional.poisson_nll_loss",
    "torch.nn.functional.cosine_embedding_loss",
    "torch.nn.functional.cross_entropy",
    "torch.nn.functional.ctc_loss",
    "torch.nn.functional.log_softmax",
    "torch.nn.functional.hinge_embedding_loss",
    "torch.nn.functional.kl_div",
    "torch.nn.functional.l1_loss",
    "torch.nn.functional.mse_loss",
    "torch.nn.functional.margin_ranking_loss",
    "torch.nn.functional.multilabel_margin_loss",
    "torch.nn.functional.multilabel_soft_margin_loss",
    "torch.nn.functional.multi_margin_loss",
    "torch.nn.functional.nll_loss",
    "torch.nn.functional.smooth_l1_loss",
    "torch.nn.functional.soft_margin_loss",
    "torch.nn.functional.triplet_margin_loss",
    "torch.nn.functional.pixel_shuffle",
    "torch.nn.functional.pixel_shuffle",
    "torch.nn.functional.pad",
    "torch.nn.functional.interpolate",
    "torch.nn.functional.upsample",
    "torch.nn.functional.interpolate",  # . This is equivalent with nn.functional.interpolate
    "torch.nn.functional.upsample_nearest",
    "torch.nn.functional.interpolate",  # This is equivalent with nn.functional.interpolate
    "torch.nn.functional.upsample_bilinear",
    "torch.nn.functional.interpolate",  # This is equivalent with nn.functional.interpolate
    "torch.nn.functional.grid_sample",
    "torch.nn.functional.affine_grid"
]
# 设置别名 F = nn.functional = torch.nn.functional
F_LIST += ["F." + name[len("torch.nn.functional."):] for name in F_LIST] + \
          [name[len("torch."):] for name in F_LIST]
F_SUPPORTED = [x for x in F_LIST if x in ALL_MAPPING]
F_UNSUPPORTED = [x for x in F_LIST if x not in ALL_MAPPING]

TORCH_DOT_LIST = [
    "torch.is_tensor",
    "torch.is_storage",
    "torch.is_complex",
    "torch.is_floating_point",
    "torch.set_default_dtype",
    "torch.get_default_dtype",
    "torch.set_default_tensor_type",
    "torch.numel",
    "torch.set_printoptions",
    "torch.set_flush_denormal",
    "torch.tensor",
    "torch.sparse_coo_tensor",
    "torch.as_tensor",
    "torch.as_strided",
    "torch.from_numpy",
    "torch.zeros",
    "torch.zeros_like",
    "torch.ones",
    "torch.ones_like",
    "torch.arange",
    "torch.range",
    "torch.linspace",
    "torch.logspace",
    "torch.eye",
    "torch.empty",
    "torch.empty_like",
    "torch.empty_strided",
    "torch.full",
    "torch.full_like",
    "torch.quantize_per_tensor",
    "torch.quantize_per_channel",
    "torch.cat",
    "torch.chunk",
    "torch.gather",
    "torch.index_select",
    "torch.masked_select",
    "torch.narrow",
    "torch.nonzero",
    "torch.reshape",
    "torch.split",
    "torch.squeeze",
    "torch.stack",
    "torch.t",
    "torch.take",
    "torch.transpose",
    "torch.unbind",
    "torch.unsqueeze",
    "torch.where",
    "torch._C.Generator",
    "torch._C.Generator.device",
    "torch._C.Generator.get_state",
    "torch._C.Generator.initial_seed",
    "torch._C.Generator.manual_seed",
    "torch._C.Generator.seed",
    "torch._C.Generator.set_state",
    "torch.seed",
    "torch.manual_seed",
    "torch.initial_seed",
    "torch.get_rng_state",
    "torch.set_rng_state",
    "torch.default_generator",
    "torch.bernoulli",
    "torch.multinomial",
    "torch.normal",
    "torch.poisson",
    "torch.rand",
    "torch.rand_like",
    "torch.randint",
    "torch.randint_like",
    "torch.randn",
    "torch.randn_like",
    "torch.randperm",
    "torch.quasirandom.SobolEngine",
    "torch.quasirandom.SobolEngine.draw",
    "torch.quasirandom.SobolEngine.fast_forward",
    "torch.quasirandom.SobolEngine.reset",
    "torch.save",
    "torch.load",
    "torch.get_num_threads",
    "torch.set_num_threads",
    "torch.get_num_interop_threads",
    "torch.set_num_interop_threads",
    "torch.no_grad",
    "torch.enable_grad",
    "torch.set_grad_enabled",
    "torch.abs",
    "torch.acos",
    "torch.add",
    "torch.addcdiv",
    "torch.addcmul",
    "torch.angle",
    "torch.asin",
    "torch.atan",
    "torch.atan2",
    "torch.bitwise_not",
    "torch.bitwise_and",
    "torch.bitwise_or",
    "torch.bitwise_xor",
    "torch.ceil",
    "torch.clamp",
    "torch.conj",
    "torch.cos",
    "torch.cosh",
    "torch.div",
    "torch.digamma",
    "torch.erf",
    "torch.erfc",
    "torch.erfinv",
    "torch.exp",
    "torch.expm1",
    "torch.floor",
    "torch.floor_divide",
    "torch.fmod",
    "torch.frac",
    "torch.imag",
    "torch.lerp",
    "torch.lgamma",
    "torch.log",
    "torch.log10",
    "torch.log1p",
    "torch.log2",
    "torch.logical_and",
    "torch.logical_not",
    "torch.logical_or",
    "torch.logical_xor",
    "torch.mul",
    "torch.mvlgamma",
    "torch.neg",
    "torch.polygamma",
    "torch.pow",
    "torch.real",
    "torch.reciprocal",
    "torch.remainder",
    "torch.round",
    "torch.rsqrt",
    "torch.sigmoid",
    "torch.sign",
    "torch.sin",
    "torch.sinh",
    "torch.sqrt",
    "torch.square",
    "torch.tan",
    "torch.tanh",
    "torch.true_divide",
    "torch.trunc",
    "torch.argmax",
    "torch.argmin",
    "torch.dist",
    "torch.logsumexp",
    "torch.mean",
    "torch.median",
    "torch.mode",
    "torch.norm",
    "torch.prod",
    "torch.std",
    "torch.std_mean",
    "torch.sum",
    "torch.unique",
    "torch.unique_consecutive",
    "torch.var",
    "torch.var_mean",
    "torch.allclose",
    "torch.argsort",
    "torch.eq",
    "torch.equal",
    "torch.ge",
    "torch.gt",
    "torch.isfinite",
    "torch.isinf",
    "torch.isnan",
    "torch.kthvalue",
    "torch.le",
    "torch.lt",
    "torch.max",
    "torch.min",
    "torch.ne",
    "torch.sort",
    "torch.topk",
    "torch.fft",
    "torch.ifft",
    "torch.rfft",
    "torch.irfft",
    "torch.stft",
    "torch.bartlett_window",
    "torch.blackman_window",
    "torch.hamming_window",
    "torch.hann_window",
    "torch.bincount",
    "torch.broadcast_tensors",
    "torch.cartesian_prod",
    "torch.cdist",
    "torch.combinations",
    "torch.cross",
    "torch.cummax",
    "torch.cummin",
    "torch.cumprod",
    "torch.cumsum",
    "torch.diag",
    "torch.diag_embed",
    "torch.diagflat",
    "torch.diagonal",
    "torch.einsum",
    "torch.flatten",
    "torch.flip",
    "torch.rot90",
    "torch.histc",
    "torch.meshgrid",
    "torch.renorm",
    "torch.repeat_interleave",
    "torch.roll",
    "torch.tensordot",
    "torch.trace",
    "torch.tril",
    "torch.tril_indices",
    "torch.triu",
    "torch.triu_indices",
    "torch.addbmm",
    "torch.addmm",
    "torch.addmv",
    "torch.addr",
    "torch.baddbmm",
    "torch.bmm",
    "torch.chain_matmul",
    "torch.cholesky",
    "torch.cholesky_inverse",
    "torch.cholesky_solve",
    "torch.dot",
    "torch.eig",
    "torch.geqrf",
    "torch.ger",
    "torch.inverse",
    "torch.det",
    "torch.logdet",
    "torch.slogdet",
    "torch.lstsq",
    "torch.lu",
    "torch.lu_solve",
    "torch.lu_unpack",
    "torch.matmul",
    "torch.matrix_power",
    "torch.matrix_rank",
    "torch.mm",
    "torch.mv",
    "torch.orgqr",
    "torch.ormqr",
    "torch.pinverse",
    "torch.qr",
    "torch.solve",
    "torch.svd",
    "torch.svd_lowrank",
    "torch.pca_lowrank",
    "torch.symeig",
    "torch.lobpcg",
    "torch.trapz",
    "torch.triangular_solve",
    "torch.compiled_with_cxx11_abi",
    "torch.result_type",
    "torch.can_cast",
    "torch.promote_types"
]
TORCH_DOT_SUPPORTED = [x for x in TORCH_DOT_LIST if x in ALL_MAPPING]
TORCH_DOT_UNSUPPORTED = [x for x in TORCH_DOT_LIST if x not in ALL_MAPPING]

TENSOR_DOT_LIST = [
    '.new_tensor',
    '.new_full',
    '.new_empty',
    '.new_ones',
    '.new_zeros',
    '.abs',
    '.abs_',
    '.acos',
    '.acos_',
    '.add',
    '.add_',
    '.addbmm',
    '.addbmm_',
    '.addcdiv',
    '.addcdiv_',
    '.addcmul',
    '.addcmul_',
    '.addmm',
    '.addmm_',
    '.addmv',
    '.addmv_',
    '.addr',
    '.addr_',
    '.allclose',
    '.angle',
    '.apply_',
    '.argmax',
    '.argmin',
    '.argsort',
    '.asin',
    '.asin_',
    '.as_strided',
    '.atan',
    '.atan2',
    '.atan2_',
    '.atan_',
    '.baddbmm',
    '.baddbmm_',
    '.bernoulli',
    '.bernoulli_',
    '.bfloat16',
    '.bincount',
    '.bitwise_not',
    '.bitwise_not_',
    '.bitwise_and',
    '.bitwise_and_',
    '.bitwise_or',
    '.bitwise_or_',
    '.bitwise_xor',
    '.bitwise_xor_',
    '.bmm',
    '.bool',
    '.byte',
    '.cauchy_',
    '.ceil',
    '.ceil_',
    '.char',
    '.cholesky',
    '.cholesky_inverse',
    '.cholesky_solve',
    '.chunk',
    '.clamp',
    '.clamp_',
    '.clone',
    '.contiguous',
    '.copy_',
    '.conj',
    '.cos',
    '.cos_',
    '.cosh',
    '.cosh_',
    '.cpu',
    '.cross',
    '.cuda',
    '.cummax',
    '.cummin',
    '.cumprod',
    '.cumsum',
    '.data_ptr',
    '.dequantize',
    '.det',
    '.dense_dim',
    '.diag',
    '.diag_embed',
    '.diagflat',
    '.diagonal',
    '.fill_diagonal_',
    '.digamma',
    '.digamma_',
    '.dim',
    '.dist',
    '.div',
    '.div_',
    '.dot',
    '.double',
    '.eig',
    '.element_size',
    '.eq',
    '.eq_',
    '.equal',
    '.erf',
    '.erf_',
    '.erfc',
    '.erfc_',
    '.erfinv',
    '.erfinv_',
    '.exp',
    '.exp_',
    '.expm1',
    '.expm1_',
    '.expand',
    '.expand_as',
    '.exponential_',
    '.fft',
    '.fill_',
    '.flatten',
    '.flip',
    '.float',
    '.floor',
    '.floor_',
    '.floor_divide',
    '.floor_divide_',
    '.fmod',
    '.fmod_',
    '.frac',
    '.frac_',
    '.gather',
    '.ge',
    '.ge_',
    '.geometric_',
    '.geqrf',
    '.ger',
    '.get_device',
    '.gt',
    '.gt_',
    '.half',
    '.hardshrink',
    '.histc',
    '.ifft',
    '.index_add_',
    '.index_add',
    '.index_copy_',
    '.index_copy',
    '.index_fill_',
    '.index_fill',
    '.index_put_',
    '.index_put',
    '.index_select',
    '.indices',
    '.int',
    '.int_repr',
    '.inverse',
    '.irfft',
    '.is_contiguous',
    '.is_complex',
    '.is_floating_point',
    '.is_pinned',
    '.is_set_to',
    '.is_shared',
    '.is_signed',
    '.item',
    '.kthvalue',
    '.le',
    '.le_',
    '.lerp',
    '.lerp_',
    '.lgamma',
    '.lgamma_',
    '.log',
    '.log_',
    '.logdet',
    '.log10',
    '.log10_',
    '.log1p',
    '.log1p_',
    '.log2',
    '.log2_',
    '.log_normal_',
    '.logsumexp',
    '.logical_and',
    '.logical_and_',
    '.logical_not',
    '.logical_not_',
    '.logical_or',
    '.logical_or_',
    '.logical_xor',
    '.logical_xor_',
    '.long',
    '.lstsq',
    '.lt',
    '.lt_',
    '.lu',
    '.lu_solve',
    '.map_',
    '.masked_scatter_',
    '.masked_scatter',
    '.masked_fill_',
    '.masked_fill',
    '.masked_select',
    '.matmul',
    '.matrix_power',
    '.max',
    '.mean',
    '.median',
    '.min',
    '.mm',
    '.mode',
    '.mul',
    '.mul_',
    '.multinomial',
    '.mv',
    '.mvlgamma',
    '.mvlgamma_',
    '.narrow',
    '.narrow_copy',
    '.ndimension',
    '.ne',
    '.ne_',
    '.neg',
    '.neg_',
    '.nelement',
    '.nonzero',
    '.norm',
    '.normal_',
    '.numel',
    '.numpy',
    '.orgqr',
    '.ormqr',
    '.permute',
    '.pin_memory',
    '.pinverse',
    '.polygamma',
    '.polygamma_',
    '.pow',
    '.pow_',
    '.prod',
    '.put_',
    '.qr',
    '.qscheme',
    '.q_scale',
    '.q_zero_point',
    '.q_per_channel_scales',
    '.q_per_channel_zero_points',
    '.q_per_channel_axis',
    '.random_',
    '.reciprocal',
    '.reciprocal_',
    '.record_stream',
    '.remainder',
    '.remainder_',
    '.renorm',
    '.renorm_',
    '.repeat',
    '.repeat_interleave',
    '.requires_grad_',
    '.reshape',
    '.reshape_as',
    '.resize_',
    '.resize_as_',
    '.rfft',
    '.roll',
    '.rot90',
    '.round',
    '.round_',
    '.rsqrt',
    '.rsqrt_',
    '.scatter',
    '.scatter_',
    '.scatter_add_',
    '.scatter_add',
    '.select',
    '.set_',
    '.share_memory_',
    '.short',
    '.sigmoid',
    '.sigmoid_',
    '.sign',
    '.sign_',
    '.sin',
    '.sin_',
    '.sinh',
    '.sinh_',
    '.size',
    '.slogdet',
    '.solve',
    '.sort',
    '.split',
    '.sparse_mask',
    '.sparse_dim',
    '.sqrt',
    '.sqrt_',
    '.square',
    '.square_',
    '.squeeze',
    '.squeeze_',
    '.std',
    '.stft',
    '.storage',
    '.storage_offset',
    '.storage_type',
    '.stride',
    '.sub',
    '.sub_',
    '.sum',
    '.sum_to_size',
    '.svd',
    '.symeig',
    '.t',
    '.t_',
    '.to',
    '.to_mkldnn',
    '.take',
    '.tan',
    '.tan_',
    '.tanh',
    '.tanh_',
    '.tolist',
    '.topk',
    '.to_sparse',
    '.trace',
    '.transpose',
    '.transpose_',
    '.triangular_solve',
    '.tril',
    '.tril_',
    '.triu',
    '.triu_',
    '.true_divide',
    '.true_divide_',
    '.trunc',
    '.trunc_',
    '.type',
    '.type_as',
    '.unbind',
    '.unfold',
    '.uniform_',
    '.unique',
    '.unique_consecutive',
    '.unsqueeze',
    '.unsqueeze_',
    '.values',
    '.var',
    '.view',
    '.view_as',
    '.where',
    '.zero_'
]
TENSOR_DOT_SUPPORTED = [x for x in TENSOR_DOT_LIST if x in ALL_MAPPING]
TENSOR_DOT_UNSUPPORTED = [x for x in TENSOR_DOT_LIST if x not in ALL_MAPPING]

ALL_2P_LIST = F_LIST + TORCH_DOT_LIST + TENSOR_DOT_LIST
ALL_TORCH_APIS = NN_LIST + F_LIST + TORCH_DOT_LIST + TENSOR_DOT_LIST
ALL_SUPPORTED = NN_SUPPORTED + F_SUPPORTED + TORCH_DOT_SUPPORTED + TENSOR_DOT_SUPPORTED
ALL_UNSUPPORTED = NN_UNSUPPORTED + F_UNSUPPORTED + TORCH_DOT_UNSUPPORTED + TENSOR_DOT_UNSUPPORTED

# TODO 使用别名
UNSUPPORTED_WARN_INFOS = {
    "nn.AdaptiveAvgPool2d": "maybe could convert to P.ReduceMean",
    "F.adaptive_avg_pool2d": "maybe could convert to P.ReduceMean",
    "F.dropout": "please use nn.Dropout in __init__()",
}
